#include <cusp/linear_operator.h>
#include <cusp/transpose.h>
#include <cusp/gallery/eye.h>
#include <cusp/array1d.h>
#include <cusp/copy.h>
#include <cusp/print.h>
#if defined(CUDA_3_2)
#include "cublas.h"
#else
#include "cublas_v2.h"
#endif

// Version of QR with Householder reflectors
// The output is Q (unitary) R (upper triangular)


namespace cusp
{

// ASSUMPTION:  All 2D arrays have column-major ordering!!!!
template <typename LinearOperator,
          typename Matrix1,
          typename Matrix2>
__host__ 
int qrhhb(LinearOperator& A,
        Matrix1& Q,
        Matrix2& R,
        size_t r)
{

  CUSP_PROFILE_SCOPED();

  typedef typename LinearOperator::value_type   ValueType;
  typedef typename LinearOperator::memory_space MemorySpace;

  typedef typename cusp::array1d<ValueType, MemorySpace>::iterator Iterator;

  size_t m = A.num_rows;
  size_t n = A.num_cols;  
  size_t blocks = n / r;       // Assumption:  n is a multiple of r

  cusp::array2d<ValueType, MemorySpace, cusp::column_major> W(m,r,0);  // H = I + WY^T
  cusp::array2d<ValueType, MemorySpace, cusp::column_major> Y(m,r,0); 
  cusp::array2d<ValueType, MemorySpace, cusp::column_major> Z(m,m,0);  // Z = WY^T
  cusp::array2d<ValueType, MemorySpace, cusp::column_major> J(m,n,0);

  R  = A;
  size_t ldr  = R.pitch;         // should be >= m, depending on padding
  size_t ldy  = Y.pitch;
  size_t ldw  = W.pitch;
  size_t ldz  = Z.pitch;

  cusp::gallery::eye(Q, m, m);
  size_t ldq  = Q.pitch;

#if !defined(CUDA_3_2)
  float alpha, beta;
  cublasStatus_t stat;
  cublasHandle_t handle;
  stat = cublasCreate(&handle);
  if (stat != CUBLAS_STATUS_SUCCESS) {
     printf ("CUBLAS initialization failed\n");
     return CUBLAS_STATUS_INTERNAL_ERROR;
  }
#endif

  for (size_t k = 0; k < blocks; ++k)
  {
    size_t s = k*r;   // Global column offset

    for (size_t j = 0; j < r; ++j)
    {
      
      size_t u = s + j;      // Global column index (zero-based)
      cusp::array1d<ValueType, MemorySpace> v(m,0);
      cusp::array1d<ValueType, MemorySpace> y(m-u,0);

      // pitch is generally = num_cols, but it may be contain padding
      thrust::copy(R.values.begin() + u*ldr+u, R.values.begin()+u*ldr+m, v.begin()+u );

      ValueType nv = cusp::blas::nrm2(v);

      if (nv == 0) { cusp::print(v); return -1; }  // null vector; algorithm must terminate softly
      if (v[u] < 0) 
	v[u] -= nv;
      else
        v[u] += nv;
      nv = cusp::blas::nrm2(v);

      // cusp::blas::scal(v.begin()+u,v.begin()+m,1/nv);
      cusp::blas::scal(v,1/nv);
      // cusp::print(v);

      thrust::copy(v.begin(), v.end(), Y.values.begin()+j*ldy );  // Copy to j-th column, from u-th row
      
      // Perform the update in the current column block:  A(u:m, u:s+r-1) = A(u:m, u:s+r-1)-2*v*v'*A(u:m,u:s+r-1)
      // Or     Sgemv(A(u:m, u:s+r-1), 1, v, 0, y);   Sger( A, -2, v, y)

#if defined(CUDA_3_2)
      cublasSgemv('t', m-u, n-u, ((float)1.0), thrust::raw_pointer_cast(&R.values[u+u*ldr]), ldr, thrust::raw_pointer_cast(&v[u]), 1, ((float)0.0), thrust::raw_pointer_cast(&y[0]), 1);
      cublasSger(m-u, r-j, -((float)2.0), thrust::raw_pointer_cast(&v[u]), 1, thrust::raw_pointer_cast(&y[0]), 1, thrust::raw_pointer_cast(&R.values[u+u*ldr]), ldr );
#else
      alpha = 1.0f;
      beta  = 0.0f;
      stat = cublasSgemv(handle, CUBLAS_OP_T, m-u, n-u, &alpha, thrust::raw_pointer_cast(&R.values[u+u*ldr]), ldr, thrust::raw_pointer_cast(&v[u]), 1, &beta, thrust::raw_pointer_cast(&y[0]), 1);
      alpha = -2.0f;
      stat = cublasSger(handle, m-u, r-j, &alpha, thrust::raw_pointer_cast(&v[u]), 1, thrust::raw_pointer_cast(&y[0]), 1, thrust::raw_pointer_cast(&R.values[u+u*ldr]), ldr );
#endif
    }

    // W = -2*Y(:,1);
    cusp::array1d<ValueType, MemorySpace> z(m-s,0);
    for (size_t j = 0; j < r; ++j) 
    {  //  W(:,j) =  -2*Y(:,j) - 2*W(:,1:j-1)*Y(:,1:j-1)'*Y(:,j)   -- uncautious reference to W.values[(0-1)*ldy] should work OK
       thrust::copy(Y.values.begin()+j*ldy, Y.values.begin()+j*ldy+m, W.values.begin()+j*ldw);
#if defined(CUDA_3_2)
       if ( j > 0 ) 
          cublasSgemv('t', m-s, r-j, ((float)1.0), thrust::raw_pointer_cast(&Y.values[(j-1)*ldy+s]), ldy, thrust::raw_pointer_cast(&Y.values[j*ldy+s]), 1, ((float)0.0), thrust::raw_pointer_cast(&z[0]), 1);
       cublasSgemv('n', m-s, r-j, -((float)2.0), thrust::raw_pointer_cast(&W.values[(j-1)*ldw+s]), ldw, thrust::raw_pointer_cast(&z[0]), 1, -((float)2.0), thrust::raw_pointer_cast(&W.values[j*ldw+s]), 1);
#else
       if ( j > 0 ) {
          alpha = 1.0f;
          beta  = 0.0f;
          stat = cublasSgemv(handle, CUBLAS_OP_T, m-s, r-j, &alpha, thrust::raw_pointer_cast(&Y.values[(j-1)*ldy+s]), ldy, thrust::raw_pointer_cast(&Y.values[j*ldy+s]), 1, &beta, thrust::raw_pointer_cast(&z[0]), 1);
       }
       alpha = -2.0f;
       beta  = -2.0f;
       stat = cublasSgemv(handle, CUBLAS_OP_N, m-s, r-j, &alpha, thrust::raw_pointer_cast(&W.values[(j-1)*ldw+s]), ldw, thrust::raw_pointer_cast(&z[0]), 1, &beta, thrust::raw_pointer_cast(&W.values[j*ldw+s]), 1);
#endif
    }

    // A(s:m,s+r:n) = A(s:m,s+r:n) + Y(s:m,:)*W(s:m,:)'*A(s:m,s+r:n);
    // First calculate Z = WY^T
#if defined(CUDA_3_2)
    cublasSgemm('n', 't', m-s, m-s, r, ((float)1.0), thrust::raw_pointer_cast(&W.values[s]), ldw, thrust::raw_pointer_cast(&Y.values[s]), ldy, ((float)0.0), thrust::raw_pointer_cast(&Z.values[s+s*ldz]), ldz);
    // A(s:m,s+r:n) = A(s:m,s+r:n) + Z'(s:m,s:m)*A(s:m,s+r:n);
    cublasSgemm('t', 'n', m-s, n-(s+r), m-s, ((float)1.0), thrust::raw_pointer_cast(&Z.values[s+s*ldz]), ldz, thrust::raw_pointer_cast(&R.values[s+(s+r)*ldr]), ldr, ((float)1.0), thrust::raw_pointer_cast(&R.values[s+(s+r)*ldr]), ldr);
    // Q(1:m,s:m) = Q(1:m,s:m)+Q(1:m,s:m)*Z(s:m,s:m)
    cublasSgemm('n', 'n', m, m-s, m-s, ((float)1.0), thrust::raw_pointer_cast(&Q.values[s*ldq]), ldq, thrust::raw_pointer_cast(&Z.values[s+s*ldz]), ldr, ((float)1.0), thrust::raw_pointer_cast(&Q.values[s*ldq]), ldq);
#else
    alpha = 1.0f;
    beta  = 0.0f; 
   stat = cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m-s, m-s, r, &alpha, thrust::raw_pointer_cast(&W.values[s]), ldw, thrust::raw_pointer_cast(&Y.values[s]), ldy, &beta, thrust::raw_pointer_cast(&Z.values[s+s*ldz]), ldz);
    // A(s:m,s+r:n) = A(s:m,s+r:n) + Z'(s:m,s:m)*A(s:m,s+r:n);
    beta  = 1.0f;
    stat = cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m-s, n-(s+r), m-s, &alpha, thrust::raw_pointer_cast(&Z.values[s+s*ldz]), ldz, thrust::raw_pointer_cast(&R.values[s+(s+r)*ldr]), ldr, &beta, thrust::raw_pointer_cast(&R.values[s+(s+r)*ldr]), ldr);
    // Q(1:m,s:m) = Q(1:m,s:m)+Q(1:m,s:m)*Z(s:m,s:m)
    stat = cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, m-s, m-s, &alpha, thrust::raw_pointer_cast(&Q.values[s*ldq]), ldq, thrust::raw_pointer_cast(&Z.values[s+s*ldz]), ldz, &beta, thrust::raw_pointer_cast(&Q.values[s*ldq]), ldq);
#endif
    
  }
  CUSP_PROFILE_DUMP();
  return 0;
}

} // end namespace cusp

